package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/pprof"
	"os"
	"strings"
)

func OkFunc(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	fmt.Fprint(w, "ok")
}

func RootFunc(w http.ResponseWriter, r *http.Request) {
	for key, vals := range r.Header {
		for _, val := range vals {
			w.Header().Add(key, val)
		}
	}
	ver, isPresent := os.LookupEnv("VERSION")
	if isPresent {
		w.Header().Set("Version", ver)
	}
	clientIP := getCurrentIP(r)
	w.WriteHeader(200)
	respMsg := fmt.Sprintf("Success! Client IP is: %s", clientIP)
	log.Printf("Success! Client IP is: %s", clientIP)
	w.Write([]byte(respMsg))
}

func getCurrentIP(r *http.Request) string {
	xForwardedFor := r.Header.Get("X-Forwarded-For")
	ip := strings.Split(xForwardedFor, ",")[0]
	if ip != "" {
		return ip
	}
	ip = r.Header.Get("X-Real-IP")
	if ip == "" {
		ip = strings.Split(r.RemoteAddr, ":")[0]
	}
	return ip
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/debug/pprof/", pprof.Index)
	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
	mux.HandleFunc("/", RootFunc)
	mux.HandleFunc("/healthz", OkFunc)
        listenPort := os.Getenv("HTTP_PORT")
	if listenPort == "" {
		listenPort = "8080"
	}
	err := http.ListenAndServe(":"+listenPort, mux)
	if err != nil {
		log.Fatalf("start http server failed, error: %s\n", err.Error())
	}
}
